package main

import (
	"encoding/binary"
	"flag"
	"fmt"
	"log"
	"net"
	"os"
	"os/signal"
	"sync"
	"syscall"
	"time"

	"github.com/google/gopacket"
	"github.com/google/gopacket/layers"
	"github.com/google/gopacket/pcap"
)

// OSPF 常量定义
const (
	OSPFHello        = 1
	OSPFDatabaseDesc = 2
	OSPFLinkStateReq = 3
	OSPFLinkStateUpd = 4
	OSPFLinkStateAck = 5

	OSPFv2HeaderLength   = 24
	HelloPacketMinLength = 20
)

// 注册自定义层类型
var (
	LayerTypeOSPFv2 = gopacket.RegisterLayerType(
		2001,
		gopacket.LayerTypeMetadata{
			Name:    "OSPFv2",
			Decoder: gopacket.DecodeFunc(decodeOSPFv2),
		},
	)
	LayerTypeOSPFv2Hello = gopacket.RegisterLayerType(
		2002,
		gopacket.LayerTypeMetadata{
			Name:    "OSPFv2Hello",
			Decoder: gopacket.DecodeFunc(decodeOSPFv2Hello),
		},
	)
)

// OSPFv2 结构体定义
type OSPFv2 struct {
	layers.BaseLayer
	Version        uint8
	Type           uint8
	PacketLength   uint16
	RouterID       uint32
	AreaID         uint32
	Checksum       uint16
	AuType         uint16
	Authentication uint64
	Contents       []byte
	Payload        []byte
}

// OSPFv2Hello 结构体定义
type OSPFv2Hello struct {
	layers.BaseLayer
	NetworkMask   uint32
	HelloInterval uint16
	Options       uint8
	Priority      uint8
	DeadInterval  uint32
	DR            uint32
	BDR           uint32
	Neighbors     []uint32
	Contents      []byte
	Payload       []byte
}

// OSPFServer 结构体定义
type OSPFServer struct {
	handle     *pcap.Handle
	localIP    net.IP
	routerID   uint32
	ifaceName  string
	stopChan   chan struct{}
	wg         sync.WaitGroup
	outputFile *os.File
	fileWriter *BufferedFileWriter
	mu         sync.Mutex
	stats      PacketStats
}

// ServerConfig 服务器配置
type ServerConfig struct {
	IfaceName   string
	LocalIP     string
	OutputPath  string
	BufferSize  int
	PromiscMode bool
	ReadTimeout time.Duration
}

// PacketStats 数据包统计
type PacketStats struct {
	TotalPackets uint64
	OSPFPackets  uint64
	HelloPackets uint64
	LastReceived time.Time
	mu           sync.RWMutex
}

// DecodeFromBytes 实现解码方法
func (o *OSPFv2) DecodeFromBytes(data []byte, df gopacket.DecodeFeedback) error {
	if len(data) < OSPFv2HeaderLength {
		return fmt.Errorf("OSPF packet too short")
	}

	o.Version = data[0]
	o.Type = data[1]
	o.PacketLength = binary.BigEndian.Uint16(data[2:4])
	o.RouterID = binary.BigEndian.Uint32(data[4:8])
	o.AreaID = binary.BigEndian.Uint32(data[8:12])
	o.Checksum = binary.BigEndian.Uint16(data[12:14])
	o.AuType = binary.BigEndian.Uint16(data[14:16])
	o.Authentication = binary.BigEndian.Uint64(data[16:24])

	o.Contents = data[:OSPFv2HeaderLength]
	o.Payload = data[OSPFv2HeaderLength:]

	return nil
}

func (h *OSPFv2Hello) DecodeFromBytes(data []byte, df gopacket.DecodeFeedback) error {
	if len(data) < HelloPacketMinLength {
		return fmt.Errorf("OSPF Hello packet too short")
	}

	h.NetworkMask = binary.BigEndian.Uint32(data[0:4])
	h.HelloInterval = binary.BigEndian.Uint16(data[4:6])
	h.Options = data[6]
	h.Priority = data[7]
	h.DeadInterval = binary.BigEndian.Uint32(data[8:12])
	h.DR = binary.BigEndian.Uint32(data[12:16])
	h.BDR = binary.BigEndian.Uint32(data[16:20])

	remainingBytes := len(data) - HelloPacketMinLength
	if remainingBytes > 0 && remainingBytes%4 == 0 {
		neighborCount := remainingBytes / 4
		h.Neighbors = make([]uint32, neighborCount)
		for i := 0; i < neighborCount; i++ {
			offset := HelloPacketMinLength + (i * 4)
			h.Neighbors[i] = binary.BigEndian.Uint32(data[offset : offset+4])
		}
	}

	h.Contents = data
	h.Payload = nil

	return nil
}

// LayerType 实现
func (o *OSPFv2) LayerType() gopacket.LayerType      { return LayerTypeOSPFv2 }
func (h *OSPFv2Hello) LayerType() gopacket.LayerType { return LayerTypeOSPFv2Hello }

// 解码函数
func decodeOSPFv2(data []byte, p gopacket.PacketBuilder) error {
	ospf := &OSPFv2{}
	err := ospf.DecodeFromBytes(data, p)
	if err != nil {
		return err
	}
	p.AddLayer(ospf)
	return p.NextDecoder(gopacket.LayerTypePayload)
}

func decodeOSPFv2Hello(data []byte, p gopacket.PacketBuilder) error {
	hello := &OSPFv2Hello{}
	err := hello.DecodeFromBytes(data, p)
	if err != nil {
		return err
	}
	p.AddLayer(hello)
	return nil
}

// NewOSPFServer 创建新的OSPF服务器
func NewOSPFServer(config ServerConfig) (*OSPFServer, error) {
	if config.IfaceName == "" {
		return nil, fmt.Errorf("interface name cannot be empty")
	}
	if config.LocalIP == "" {
		return nil, fmt.Errorf("local IP cannot be empty")
	}
	if config.OutputPath == "" {
		return nil, fmt.Errorf("output path cannot be empty")
	}
	if config.BufferSize <= 0 {
		config.BufferSize = 4096
	}
	if config.ReadTimeout == 0 {
		config.ReadTimeout = pcap.BlockForever //永远阻塞
	}

	handle, err := pcap.OpenLive(
		config.IfaceName,
		65535,
		config.PromiscMode,
		config.ReadTimeout,
	)
	if err != nil {
		return nil, fmt.Errorf("failed to open interface %s: %v", config.IfaceName, err)
	}

	err = handle.SetBPFFilter("proto ospf")
	if err != nil {
		handle.Close()
		return nil, fmt.Errorf("failed to set BPF filter: %v", err)
	}

	file, err := os.OpenFile(config.OutputPath, os.O_CREATE|os.O_APPEND|os.O_WRONLY, 0644)
	if err != nil {
		handle.Close()
		return nil, fmt.Errorf("failed to open output file %s: %v", config.OutputPath, err)
	}

	bufferedWriter := NewBufferedFileWriter(file, config.BufferSize)

	ipAddr := net.ParseIP(config.LocalIP)
	if ipAddr == nil {
		handle.Close()
		file.Close()
		return nil, fmt.Errorf("invalid IP address: %s", config.LocalIP)
	}

	return &OSPFServer{
		handle:     handle,
		localIP:    ipAddr,
		routerID:   binary.BigEndian.Uint32(ipAddr.To4()),
		ifaceName:  config.IfaceName,
		stopChan:   make(chan struct{}),
		outputFile: file,
		fileWriter: bufferedWriter,
	}, nil
}

// Start 启动服务器
func (s *OSPFServer) Start() error {
	log.Printf("OSPF Server starting on interface %s...", s.ifaceName)

	s.wg.Add(1)
	go s.capturePackets()

	s.wg.Add(1)
	go s.printStats()

	return nil
}

// Stop 停止服务器
func (s *OSPFServer) Stop() {
	log.Println("Stopping OSPF server...")
	close(s.stopChan)
	s.wg.Wait()

	if s.fileWriter != nil {
		s.fileWriter.Flush()
	}
	if s.outputFile != nil {
		s.outputFile.Close()
	}
	if s.handle != nil {
		s.handle.Close()
	}
	log.Println("OSPF server stopped")
}

// capturePackets 捕获和处理OSPF数据包
func (s *OSPFServer) capturePackets() {
	defer s.wg.Done()

	packetSource := gopacket.NewPacketSource(s.handle, s.handle.LinkType())
	packetSource.NoCopy = true

	for {
		select {
		case <-s.stopChan:
			return
		default:
			packet, err := packetSource.NextPacket()
			if err != nil {
				if err == pcap.NextErrorTimeoutExpired {
					continue
				}
				log.Printf("Error reading packet: %v", err)
				continue
			}

			s.updateStats()
			if err := s.handlePacket(packet); err != nil {
				log.Printf("Error handling packet: %v", err)
			}
		}
	}
}

// handlePacket 处理单个数据包
func (s *OSPFServer) handlePacket(packet gopacket.Packet) error {
	ipLayer := packet.Layer(layers.LayerTypeIPv4)
	if ipLayer == nil {
		return fmt.Errorf("not an IPv4 packet")
	}
	ip, _ := ipLayer.(*layers.IPv4)

	if ip.Protocol != 89 {
		return fmt.Errorf("not an OSPF packet")
	}

	s.stats.mu.Lock()
	s.stats.OSPFPackets++
	s.stats.mu.Unlock()

	ospfLayer := packet.Layer(LayerTypeOSPFv2)
	if ospfLayer == nil {
		return fmt.Errorf("not a valid OSPF packet")
	}
	ospf, _ := ospfLayer.(*OSPFv2)

	if err := s.validateOSPFPacket(ospf); err != nil {
		return fmt.Errorf("invalid OSPF packet: %v", err)
	}

	if ospf.Type != OSPFHello {
		return nil
	}

	s.stats.mu.Lock()
	s.stats.HelloPackets++
	s.stats.mu.Unlock()

	applicationLayer := packet.ApplicationLayer()
	if applicationLayer == nil {
		return fmt.Errorf("no application layer found")
	}

	payload := applicationLayer.Payload()
	if len(payload) < 20 {
		return fmt.Errorf("hello packet too short")
	}

	deadInterval := binary.BigEndian.Uint32(payload[8:12])
	data := make([]byte, 4)
	binary.BigEndian.PutUint32(data, deadInterval)

	processedData, err := s.processData(data)
	if err != nil {
		return fmt.Errorf("failed to process data: %v", err)
	}

	if err := s.WriteToFile(processedData); err != nil {
		return fmt.Errorf("failed to write data: %v", err)
	}

	log.Printf("Received OSPF Hello from %v, Router ID: %v", ip.SrcIP, ospf.RouterID)
	return nil
}

// validateOSPFPacket 验证OSPF数据包
func (s *OSPFServer) validateOSPFPacket(ospf *OSPFv2) error {
	if ospf.Version != 2 {
		return fmt.Errorf("invalid OSPF version: %d", ospf.Version)
	}
	return nil
}

// processData 处理接收到的数据
func (s *OSPFServer) processData(data []byte) ([]byte, error) {
	lastNonZero := len(data) - 1
	for ; lastNonZero >= 0 && data[lastNonZero] == 0; lastNonZero-- {
	}
	data = data[:lastNonZero+1]

	if len(data) == 0 {
		return nil, fmt.Errorf("empty data after processing")
	}

	return data, nil
}

// WriteToFile 写入数据到文件
func (s *OSPFServer) WriteToFile(data []byte) error {
	s.mu.Lock()
	defer s.mu.Unlock()

	timestamp := time.Now().Format("2006-01-02 15:04:05.000")
	entry := fmt.Sprintf("[%s] Received data: %s\n", timestamp, string(data))

	return s.fileWriter.Write([]byte(entry))
}

// updateStats 更新统计信息
func (s *OSPFServer) updateStats() {
	s.stats.mu.Lock()
	defer s.stats.mu.Unlock()
	s.stats.TotalPackets++
	s.stats.LastReceived = time.Now()
}

// printStats 定期打印统计信息
func (s *OSPFServer) printStats() {
	defer s.wg.Done()
	ticker := time.NewTicker(10 * time.Second)
	defer ticker.Stop()

	for {
		select {
		case <-s.stopChan:
			return
		case <-ticker.C:
			s.stats.mu.RLock()
			log.Printf("Statistics - Total: %d, OSPF: %d, Hello: %d, Last: %v",
				s.stats.TotalPackets,
				s.stats.OSPFPackets,
				s.stats.HelloPackets,
				s.stats.LastReceived)
			s.stats.mu.RUnlock()
		}
	}
}

// BufferedFileWriter 带缓冲的文件写入器
type BufferedFileWriter struct {
	file   *os.File
	buffer []byte
	size   int
	mu     sync.Mutex
}

func NewBufferedFileWriter(file *os.File, bufferSize int) *BufferedFileWriter {
	return &BufferedFileWriter{
		file:   file,
		buffer: make([]byte, 0, bufferSize),
		size:   bufferSize,
	}
}

func (w *BufferedFileWriter) Write(data []byte) error {
	w.mu.Lock()
	defer w.mu.Unlock()

	w.buffer = append(w.buffer, data...)

	if len(w.buffer) >= w.size {
		if err := w.Flush(); err != nil {
			return err
		}
	}

	return nil
}

func (w *BufferedFileWriter) Flush() error {
	if len(w.buffer) == 0 {
		return nil
	}

	_, err := w.file.Write(w.buffer)
	if err != nil {
		return err
	}

	w.buffer = w.buffer[:0]
	return w.file.Sync()
}

// 命令行参数结构体
type CommandLineArgs struct {
	ifaceName   string
	localIP     string
	outputPath  string
	bufferSize  int
	promiscMode bool
}

// 解析命令行参数
func parseArgs() CommandLineArgs {
	args := CommandLineArgs{}

	flag.StringVar(&args.ifaceName, "i", "", "Network interface name (required)")
	flag.StringVar(&args.localIP, "ip", "", "Local IP address (required)")
	flag.StringVar(&args.outputPath, "o", "ospf_data.txt", "Output file path")
	flag.IntVar(&args.bufferSize, "b", 4096, "Buffer size for file writing")
	flag.BoolVar(&args.promiscMode, "promisc", true, "Enable promiscuous mode")

	flag.Parse()

	if args.ifaceName == "" || args.localIP == "" {
		flag.Usage()
		os.Exit(1)
	}

	return args
}

// 列出可用的网络接口
func listInterfaces() {
	devices, err := pcap.FindAllDevs()
	if err != nil {
		log.Printf("Error listing interfaces: %v", err)
		return
	}

	log.Println("Available network interfaces:")
	for _, device := range devices {
		log.Printf("Interface: %s", device.Name)
		for _, address := range device.Addresses {
			log.Printf("  IP address: %s", address.IP)
			log.Printf("  Netmask: %s", address.Netmask)
		}
	}
}

func main() {
	// 设置日志格式
	log.SetFlags(log.Ldate | log.Ltime | log.Lmicroseconds)

	// 解析命令行参数
	args := parseArgs()

	// 创建服务器配置
	config := ServerConfig{
		IfaceName:   args.ifaceName,
		LocalIP:     args.localIP,
		OutputPath:  args.outputPath,
		BufferSize:  args.bufferSize,
		PromiscMode: args.promiscMode,
		ReadTimeout: pcap.BlockForever,
	}

	// 创建服务器实例
	server, err := NewOSPFServer(config)
	if err != nil {
		log.Printf("Available interfaces:")
		listInterfaces()
		log.Fatalf("Failed to create OSPF server: %v", err)
	}

	// 启动服务器
	if err := server.Start(); err != nil {
		log.Fatalf("Failed to start OSPF server: %v", err)
	}

	// 设置信号处理
	sigChan := make(chan os.Signal, 1)
	signal.Notify(sigChan, os.Interrupt, syscall.SIGTERM)

	// 等待终止信号
	sig := <-sigChan
	log.Printf("Received signal %v, shutting down...", sig)

	// 优雅关闭
	server.Stop()
}
